"""synergynet test"""
import mindspore as ms
import numpy as np
import mindspore.dataset as ds
from mindspore import Tensor, context
from mind3d.dataset.aflw2000 import AFLW2000
from mind3d.dataset.DDFAdataset import DDFADataset
from mind3d.models.blocks.i2p import I2P
from mindspore import Tensor


def test_aflw2000():
    """
    test the function of AFLW2000
    """
    root = '/data1/hujingsong/synergynet/aflw2000_data/AFLW2000-3D_crop'
    filelists = '/data1/hujingsong/synergynet/aflw2000_data/AFLW2000-3D_crop.list'

    dataset_generator = AFLW2000(filelists=filelists, root=root, transform=True)
    output = dataset_generator.__getitem__(0)[0]
    assert output.shape == (1, 3, 120, 120), 'output not match'
    print("################################")
    print("AFLW2000 test passed!")
    print("################################")

def test_DDFADataset():
    """
    test the function of DDFADataset
    """
    context.set_context(device_id=5, mode=context.GRAPH_MODE)
    dataset_generator = DDFADataset(root='./SynergyNet/3dmm_data/train_aug_120x120',
                                    filelists='./SynergyNet/3dmm_data/train_aug_120x120.list.train',
                                    param_fp='./SynergyNet/3dmm_data/param_all_norm_v201.pkl',
                                    )
    # data = dataset_generator[0]
    # dataset = ds.GeneratorDataset(dataset_generator, ["data", "target"], shuffle=False)
    output = dataset_generator.__getitem__(0)[0]
    assert output.shape == (1, 3, 120, 120), 'output not match'
    print("################################")
    print("DDFADataset test passed!")
    print("################################")


def test_I2P():
    """
    test the function of I2P
    """
    data = Tensor(np.ones((1,3,120,120)), ms.float32)
    out = I2P(data)
    print(out.shape)
    assert out.shape == (1,62), 'output shape not match'
    print("################################")
    print("I2P test passed!")
    print("################################")


if __name__ == "__main__":
    test_aflw2000()
    test_I2P()
    test_DDFADataset()
